#!/bin/sh

# This script unlocks encrypted partitions on UEFI systems
# with secure boot enabled and encryption key protected by a TPM2.0
#
# The partition layout of such system is:
# sda         // Boot device
# |-sda1      // Unencrypted EFI partition, contains the bootloader and encrypted disk passphrase
# |-sda2      // Encrypted boot partition
# | `-dm-0    // Unlocked boot partition, contains balena-specifics, e.g. config.json
# |-sda3      // Encrypted rootA partition
# | `-dm-1    // Unlocked rootA partition
# |-sda4      // Encrypted rootB partition
# | `-dm-2    // Unlocked rootB partition
# |-sda5      // Encrypted state partition
# | `-dm-3    // Unlocked state partition
# `-sda6      // Encrypted data partition
#   `-dm-4    // Unlocked data partition
#
# When this script executes, it should only find the encrypted partitions.
# The script unlocks the partitions using the metadata stored in the EFI
# partition in combination with secrets stored in the TPM.
# As a result the dm-0..dm-5 devices are created.
#
# After successfully unlocking a partition, the default udev rules will
# in fact create two instances of the unlocked device:
# * /dev/dm-X - the "dm-X" part is referred to as "KNAME".
#   This is what the kernel uses internally to identify the device
# * /dev/mapper/luks-$UUID - the "luks-$UUID" part is referred to as "NAME".
#   This is a user-defined name, we chose "luks-$UUID", the same as Fedora.
# The two devices are identical and can be used interchangeably, but this
# script tries to prefer KNAME as it seems more practical for scripting.

# shellcheck disable=SC1091
. /usr/libexec/os-helpers-logging
. /usr/libexec/os-helpers-fs
. /usr/libexec/os-helpers-tpm2
. /usr/libexec/os-helpers-efi
. /usr/sbin/balena-config-defaults

ensure_luks2() {
    LUKS_DEVICE="$1"

    LUKS_VERSION=$(cryptsetup luksDump "${LUKS_DEVICE}" | grep "^Version:" | cut -f 2)

    if [ "${LUKS_VERSION}" = 1 ]; then
        info "Converting ${LUKS_DEVICE} to LUKS2"
        cryptsetup convert -q --type luks2 "${LUKS_DEVICE}"
    fi
}

cryptsetup_enabled() {
    # Flasher should not try to unlock the partitions
    if [ "$bootparam_flasher" = "true" ]; then
        return 1
    fi

    # Ensure that secure boot is enabled and in user mode before unlocking
    if ! user_mode_enabled; then
        info "Won't attempt to decrypt drives because secure boot is not enabled"
        return 1
    fi

    return 0
}

cryptsetup_run() {
    # Die if anything fails here
    set -e

    # Don't load the passphrase into memory if the IOMMU is disabled
    if pcr7_has_dma_disabled_event; then
        fail "The IOMMU is disabled, won't load the LUKS passphrase into memory"
    fi

    boot_part_assert

    EFI_MOUNT_DIR="/efi"
    EFI_DEV=$(get_state_path_from_label "${BALENA_NONENC_BOOT_LABEL}")
    mkdir "$EFI_MOUNT_DIR"
    mount "$EFI_DEV" "$EFI_MOUNT_DIR"

    PCRS="sha256:0,2,3,7"
    printf "Attempting to unlock LUKS passphrase with PCRS %s\n" "$PCRS"
    tpm2_pcrread "${PCRS}"

    PASSPHRASE_FILE=/balena-luks.pwd
    SESSION_CTX="$(mktemp -t)"
    POLICY_PATH="$(find "$EFI_MOUNT_DIR" -name "policies.*")"
    tpm2_startauthsession --policy-session -S "${SESSION_CTX}"
    tpm2_policypcr -S "${SESSION_CTX}" -l "${PCRS}"

    trap 'tpm2_flushcontext "${SESSION_CTX}"' EXIT

    # combined multiple policies with tpm2_policyor
    POLICIES="$(find "${POLICY_PATH}" -type f | sort | xargs)"
    if [ "$(echo "${POLICIES}" | wc -w)" -gt 1 ]; then
        tpm2_policyor -S "${SESSION_CTX}" "sha256:$(echo "${POLICIES}" | sed 's/ /,/g')"
    fi

    if tpm_nvram_retrieve_passphrase "session:${SESSION_CTX}" "${PASSPHRASE_FILE}"; then
        info "Successfully retrieved passphrase from TPM NVRAM"
    else
        umount "$EFI_MOUNT_DIR"
        fail "Failed to unlock LUKS passphrase using the TPM"
    fi

    tpm2_flushcontext "${SESSION_CTX}" >/dev/null 2>&1

    luks_parts_premount_assert

    # Unlock all the partitions - cryptsetup luksOpen does not wait for udev processing
    # of the DM device to complete, it just kicks off the process and returns.
    # Since this is async, we can perform all the luksOpens here, note the device names
    # and wait for them in a separate loop later
    LUKS_UNLOCKED=""
    BOOT_DEVICE=$(lsblk -nlo pkname "${EFI_DEV}")
    LUKS_PARTITIONS=$(lsblk -nlo "kname,uuid,fstype,partlabel" "/dev/${BOOT_DEVICE}" | grep "crypto_LUKS")
    for PART_NAME in ${DEFAULT_PARTITION_NAMES}; do
        LUKS_UUID=$(echo "${LUKS_PARTITIONS}" | grep " \(balena\|resin\)-${PART_NAME}$" | awk '{print $2}')

        if [ -z "${LUKS_UUID}" ]; then
            fail "Partition '${PART_NAME}' not found"
        fi

        if [ "$(echo "${LUKS_UUID}" | wc -l)" -gt 1 ]; then
            fail "More than one '${PART_NAME}' partition found"
        fi

        ensure_luks2 "/dev/disk/by-uuid/${LUKS_UUID}"
        cryptsetup luksOpen --key-file "${PASSPHRASE_FILE}" "UUID=${LUKS_UUID}" "luks-${LUKS_UUID}"
        LUKS_UNLOCKED="${LUKS_UNLOCKED} luks-${LUKS_UUID}"
    done

    # Wait for udev processing of each unlocked device
    for DM_NAME in ${LUKS_UNLOCKED}; do
        wait4udev "/dev/mapper/${DM_NAME}"
    done

    # Perform sanity checks after unlocking.
    # We know what the system should look like after the partitions are unlocked.
    # We want to make sure that the newly unlocked partitions are the ones
    # we are going to actually use, there is nothing missing and nothing extra.
    luks_parts_postmount_assert

    rm -f "$PASSPHRASE_FILE"
    umount "$EFI_MOUNT_DIR"
    rmdir "$EFI_MOUNT_DIR"

    # Revert dying on error
    set +e
}
